import numpy as np
import pandas as pd
import torch
import os
import torch.nn.functional as F
import cv2
# import matplotlib.pyplot as plt
# import skimage.io as io
# import matplotlib.pyplot as plt


def make_single_heatmap(img_width, img_height, c_x, c_y, sigma):
    ## 如果坐标为0，则表示该点实际为缺失点。生成空白的heatmap
    if c_x==0 or c_y ==0:
        sigma = 0
    X1 = np.linspace(1, img_width, img_width)
    Y1 = np.linspace(1, img_height, img_height)
    [X, Y] = np.meshgrid(X1, Y1)
    X = X - c_x
    Y = Y - c_y
    D2 = X * X + Y * Y
    E2 = 2.0 * sigma * sigma
    Exponent = D2 / E2
    heatmap = np.exp(-Exponent)
    return heatmap[np.newaxis,...]

def make_all_heatmap(img_width, img_height, anno, sigma=21):
    # assert len(anno) == 8,'the len of anno is not 8'
    ## 提供行列的形式作为位置信息，传入make_single_heatmap时参数转换成了坐标的形式了
    ## 如果该点标记缺失，读取的坐标是[0,0]。生成空白的heatmap
    heatmap = make_single_heatmap(img_width, img_height,anno[0][1],anno[0][0],sigma=sigma)
    for single_anno in anno[1:]:
        single_heatmap = make_single_heatmap(img_width, img_height,single_anno[1],single_anno[0],sigma=sigma)
        heatmap = np.concatenate((heatmap,single_heatmap),axis=0)
    return heatmap

def make_one_heatmap(img_width, img_height, anno, sigma=21):
    #将多个坐标点的到的高斯分布生成在同一个heatmap上。
    heatmap = make_single_heatmap(img_width, img_height,anno[0][1],anno[0][0],sigma=sigma)
    for single_anno in anno[1:]:
        single_heatmap = make_single_heatmap(img_width, img_height,single_anno[1],single_anno[0],sigma=sigma)
        heatmap = np.maximum(heatmap,single_heatmap)
    return heatmap


def make_group_heatmap(img_width, img_height, anno, sigma=[11,12,13,14]):
    ## 按照分组的方式，同一组关键的高斯核放置在同一层heatmap上。
    ## 四组分组分别设置不同的sigma
    ## 多组heatmap concat
    heatmap = make_one_heatmap(img_width, img_height, anno[0], sigma=sigma[0])
    for index,single_anno in enumerate(anno[1:]):
        single_heatmap = make_one_heatmap(img_width, img_height, single_anno, sigma=sigma[index+1])
        heatmap = np.concatenate((heatmap, single_heatmap), axis=0)
    return heatmap

def group_heatmap_to_anno(heatmap):
    ## 将heatmap还原成坐标。针对侧位片进行。sp,hc，fc单独一层heatmap，其他点一层heatmap
    def heatmap_nms(heat, kernel=15):
        pad = (kernel - 1) // 2
        hmax = F.max_pool2d(
            heat[np.newaxis, ...], (kernel, kernel), stride=1, padding=pad
        )
        ## keep中，经过maxpool操作的最大值位1,其余值为0,然后在和heat相乘，保留了最大值的数值，其他的均抑制成0
        keep = (hmax == heat).float()
        return (heat * keep)[0]
    
    def get_topk(heatmap, k=6,kernel = 15):
        ## 得到数值最大的k个点的坐标。顺序按照从上到下从左到右的顺序。
        # 可以保证前三个坐标是上面的三个节点，后三个坐标是下面的三个节点，然后再根据列坐标进行排序即可确定位置。
        anno_list = []
        heatmap = heatmap_nms(heatmap,kernel=kernel)
        heatmap_view = heatmap.view(-1)
        top_k_num = torch.topk(heatmap_view, k)[0]
        heatmap = heatmap.cpu()
        anno = np.where(heatmap >= top_k_num[-1].cpu())
        for i, j in zip(anno[0][:k], anno[1][:k]):
            ## 转换时候会出现偏差，需要+1来进行调整。
            # 可后续进行验证对比，比较+1和不+那种情况下loss较小
            anno_list.append((i + 1, j + 1))
        ## 需要对6个节点进行排序来确定位置。
        ## 分别针对上面三个节点和下面的三个节点，使用其列坐标来进行排序操作
        anno_list = sorted(anno_list[:3], key=lambda j: j[1]) + sorted(anno_list[3:], key=lambda j: j[1])
        return anno_list
    
    ## 图像上端处于肩部位置的8个节点
    anno_0 = get_topk(heatmap[0], k=8,kernel=15)
    anno_0 = sorted(anno_0, key=lambda j: j[1])
    point_14 = sorted(anno_0[6:],key=lambda j: j[0])[0]
    point_15 = sorted(anno_0[6:],key=lambda j: j[0])[1]
    point_12 = sorted(anno_0[4:6], key=lambda j: j[0])[0]
    point_13 = sorted(anno_0[4:6], key=lambda j: j[0])[1]
    point_16 = sorted(anno_0[2:4], key=lambda j: j[0])[0]
    point_17 = sorted(anno_0[2:4], key=lambda j: j[0])[1]
    point_18 = sorted(anno_0[:2], key=lambda j: j[0])[0]
    point_19 = sorted(anno_0[:2], key=lambda j: j[0])[1]
    
    ## 胸椎棘突四个节点
    anno_1 = get_topk(heatmap[1], k=4,kernel=15)
    anno_1 = sorted(anno_1, key=lambda j: j[0])
    point_0,point_1,point_2,point_3 = anno_1[0],anno_1[1],anno_1[2],anno_1[3]
    

    ## 肋骨，左右各4 个，一共8个
    anno_2 = get_topk(heatmap[2], k=8,kernel=21)
    anno_2 = sorted(anno_2, key=lambda j: j[1])
    anno_2_l = sorted(anno_2[:4], key=lambda j: j[0])
    anno_2_r = sorted(anno_2[4:], key=lambda j: j[0])
    point_4,point_5,point_6,point_7 = anno_2_l[0],anno_2_l[1],anno_2_l[2],anno_2_l[3]
    point_8,point_9,point_10,point_11 = anno_2_r[0],anno_2_r[1],anno_2_r[2],anno_2_r[3]
    
    
    ## 肺部节点，左右各3个，一共6个
    anno_3 = get_topk(heatmap[3],k=6,kernel=49)
    anno_3 = sorted(anno_3, key=lambda j: j[0])
    point_20,point_21 = anno_3[0],anno_3[1]
    anno_3 = sorted(anno_3[2:], key=lambda j: j[1])
    point_22,point_23,point_24,point_25 = anno_3[0],anno_3[3],anno_3[1],anno_3[2]
    return [point_0,point_1,point_2,point_3,point_4,point_5,point_6,point_7,point_8,point_9,point_10,point_11,point_12,point_13,point_14,point_15,
            point_16,point_17,point_18,point_19,point_20,point_21,point_22,point_23,point_24,point_25]
    
    
def group_point(anno):
    ## 将point进行分组操作。统一组存在一个list'中，然后所有的list组成一个list
    ## 图像上端处于肩部位置的8个节点。
    point_0 = anno[12:20]
    ## 胸椎棘突四个节点
    point_1 = anno[:4]
    ## 肋骨，左右各4 个，一共8个
    point_2 = anno[4:12]
    ## 肺部节点，左右各3个，一共6个
    point_3 = anno[20:]
    return [point_0,point_1,point_2,point_3]
    
    



def heatmap_to_anno(heatmap):
    heatmap = heatmap.cpu()
    anno = []
    for index in range(heatmap.shape[0]):
        ## 如果为空白的heatmap，则坐标为0
        if heatmap[index].max() == 0:
            anno.append((0,0))
        else:
            single_anno = np.where(heatmap[index] == heatmap[index].max())
            anno.append((int(single_anno[0][0]),int(single_anno[1][0])))
    return anno




if __name__ == '__main__':
    pass
